package org.example.algorithm.binarytree;

import org.example.algorithm.TreeNode;

import java.util.HashMap;
import java.util.Map;

//路径总和 III
public class PathSumSolution {

    public int pathSum(TreeNode root, int targetSum) {
        if (root == null) {
            return 0;
        }
        Map<Long, Integer> prefixSumMap = new HashMap<>();
        prefixSumMap.put(0L, 1);
        return dfs(root, targetSum, prefixSumMap, 0);
    }

    private int dfs(TreeNode root, int targetSum, Map<Long, Integer> prefixSumMap, long curSum) {
        if (root == null) {
            return 0;
        }

        curSum = curSum+root.val;

        int res = prefixSumMap.getOrDefault(curSum-targetSum, 0);
        prefixSumMap.put(curSum, prefixSumMap.getOrDefault(curSum,0) + 1);
        res += dfs(root.left, targetSum, prefixSumMap, curSum);
        res += dfs(root.right, targetSum, prefixSumMap, curSum);
        prefixSumMap.put(curSum, prefixSumMap.getOrDefault(curSum,0) - 1);
        return res;
    }

    //以每个节点为起点遍历
    public int pathSum2(TreeNode root, int targetSum) {
        if (root == null) {
            return 0;
        }
        int res = calPathSum(root, targetSum);
        res += pathSum2(root.left, targetSum);
        res += pathSum2(root.right, targetSum);
        return res;
    }

    private int calPathSum(TreeNode root, int targetSum) {
        if (root == null) {
            return 0;
        }
        if (root.val == targetSum) {
            return 1;
        }
        return calPathSum(root.left, targetSum-root.val) + calPathSum(root.right, targetSum-root.val);
    }

    public static void main(String[] args) {
        TreeNode root = new TreeNode(1);
        root.left = new TreeNode(2);
        root.right = new TreeNode(6);
        root.left.left = new TreeNode(3);
        root.left.right = new TreeNode(2);
        root.left.right.left = new TreeNode(1);
        root.left.right.right = new TreeNode(2);
        PathSumSolution solution = new PathSumSolution();
        int res = solution.pathSum(root, 6);
        System.out.println(res);
    }
}
